import asyncio
import json
import traceback
from contextlib import AsyncExitStack
from typing import Optional

from dotenv import load_dotenv
from loguru import logger
from mcp import ClientSession, StdioServerParameters, ListToolsResult
from mcp.client.stdio import stdio_client
from mcp.types import CallToolResult
from openai import OpenAI
from openai.types.chat import ChatCompletion

load_dotenv()

api_key = "123456"
# base_url = "http://172.20.50.49:10000/v1"
base_url = "http://172.20.50.49:10001/v1"
# api_key = "sk-b80acbc44b764720a59d57c62d06ea63"
# base_url = "https://api.deepseek.com"
print(f"base_url: {base_url}")

class MCPClient:
    def __init__(self):
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.deepseek = OpenAI(api_key=api_key, base_url=base_url)

    async def connect_to_server(self, server_script_path: str):
        """连接到 MCP 服务器

        Args:
            server_script_path: 服务器脚本的路径 (.py 或 .js)
        """
        is_python = server_script_path.endswith(".py")
        is_js = server_script_path.endswith(".js")
        if not is_python and not is_js:
            raise ValueError("Invalid server script path")

        command = "python" if is_python else "node"
        server_params = StdioServerParameters(
            command=command, 
            args=[server_script_path], 
            env=None
        )

        stdio_transport = await self.exit_stack.enter_async_context(
            stdio_client(server_params)
        )
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(
            ClientSession(self.stdio, self.write)
        )

        await self.session.initialize()

        response = await self.session.list_tools()

        tools = response.tools
        logger.success("成功连接服务端 - 脚本路径: {}", server_script_path)
        logger.debug("服务端支持工具列表: {}", [tool.name for tool in tools])
        logger.debug("连接参数: 命令={}, 参数={}", command, [server_script_path])

    async def process_query(self, query: str) -> str:
        response:ListToolsResult = await self.session.list_tools()
        logger.debug("服务端支持工具列表: {}", [tool.name for tool in response.tools])
        available_tools = [
            {
                "type": "function",
                "function": {
                    "name": tool.name.strip(),
                    "description": tool.description,
                    "parameters": tool.inputSchema
                }
            }
            for tool in response.tools
        ]

        messages = [
            {
                "role": "system", 
                "content": "请判断是否需要调用工具，并调用适合的工具，回复请用中文。"
            },   
            {"role": "user", "content": query},
        ]
        available_tools = [
            {
                "type": "function",
                "function": {
                    "name": tool.name.strip(),
                    "description": tool.description,
                    "parameters": tool.inputSchema
                }
            }
            for tool in response.tools
        ]
        logger.info(f"客户端构造工具定义:\n{json.dumps(available_tools, indent=2, ensure_ascii=False)}")
        # noinspection PyTypeChecker
        response: ChatCompletion = self.deepseek.chat.completions.create(
            # model="deepseek-chat",
            model="qwen:7b",
            # model="deepseek-r1:7b",
            messages=messages,
            max_tokens=1000,
            tools=available_tools,
            # tool_choice={"type": "function", "function": {"name": available_tools[0]["function"]["name"]}},
            tool_choice="auto",
            stream=False,
        )
        
        logger.info(f"HTTP请求详情:\nURL: {base_url}\nid: {response.id}")  # 新增请求日志

        
        #处理响应并处理工具调用
        final_text = []
        #assistant_message_content = []
        logger.info(f"{response}")
        for choice in response.choices:
            if hasattr(choice.message, 'tool_calls') and choice.message.tool_calls:#  处理工具调用
                logger.debug(f"[工具调用] 检测到 {len(choice.message.tool_calls)} 个工具请求")
                for tool_call in choice.message.tool_calls:
                    tool_name = tool_call.function.name
                    tool_args = tool_call.function.arguments
                    
                    # 记录原始参数和解析后的参数
                    logger.debug(f"[工具调用] 开始处理 - 工具名: {tool_name}")
                    logger.debug(f"原始参数JSON字符串:\n{tool_args}")
                    logger.debug(f"消息对象ID: {response.id}")
                    logger.debug(f"用户查询内容: {query}")
                    
                    try:
                        parsed_args = json.loads(tool_args)
                        logger.info(f"解析后参数: {parsed_args}")
                    except json.JSONDecodeError as e:
                        logger.error(f"参数解析失败: {e}")
                        parsed_args = {}

                    # 执行工具调用
                    result: CallToolResult = await self.session.call_tool(
                        tool_name, parsed_args
                    )

                    logger.debug(f"响应原始数据:\n{json.dumps([item.text for item in result.content], ensure_ascii=False, indent=2)}")

                    tool_result_text = result.content[0].text
                    final_text.append(str(tool_result_text) if tool_result_text is not None else "")
                    # 先添加assistant消息
                    messages.append({
                        "role": "assistant",
                        "content": None,
                        "tool_calls": [{
                            "id": tool_call.id,
                            "type": "function",
                            "function": {
                                "name": tool_name,
                                "arguments": tool_args
                            }
                        }]
                    })
                    
                    # 再添加工具响应
                    messages.append({
                        "role": "tool",
                        "content": tool_result_text,
                        "tool_call_id": tool_call.id,
                        "name": tool_name
                    })

                # 确保 content 是可序列化的
                if hasattr(response.choices[0].message.content, 'to_json'):
                    content = response.choices[0].message.content.to_json()
                else:
                    content = str(response.choices[0].message.content) if response.choices[0].message.content is not None else ""
                final_text.append(content)

                # Call the model again with the tool result
                # noinspection PyTypeChecker
                logger.debug(f"messages: {messages}")

                # noinspection PyTypeChecker
                response = self.deepseek.chat.completions.create(
                    model="qwen:7b",
                    # model="deepseek-chat",
                    # model="deepseek-r1:7b",
                    messages=messages,
                    max_tokens=1000,
                    tools=available_tools,
                    stream=False,
                )
                final_text.append(response.choices[0].message.content)
                #logger.info(f"{response}")
                logger.info(f"{final_text}")

        return "\n".join(final_text)

    async def chat_loop(self):
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")

        while True:
            try:
                query = input("\nYour Query: ").strip()
                if query.lower() == "quit":
                    print("Exiting...")
                    break

                response = await self.process_query(query)
                print("\n" + response)
            except Exception as e:
                print(f"Error: {e}")
                traceback.print_exc()

    async def cleanup(self):
        await self.exit_stack.aclose()
        print("MCP Client Cleaned Up!")


async def main():
    print(sys.argv)
    if len(sys.argv) < 2:
        print("Usage: python client.py <path_to_server_script>")
        sys.exit(1)

    client = MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    finally:
        await client.cleanup()


if __name__ == "__main__":
    import sys

    asyncio.run(main())

